package com.ai.model.kimi;

import com.ai.model.ChatRequest;
import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import okhttp3.*;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;

/**
 * kimi大模型调用
 *
 * @author hh
 * @since 2024/6/1
 */
@Component
public class KimiChatAiRequest implements ChatRequest {

    private final KimiApiProperties kimiApiProperties;

    public KimiChatAiRequest(KimiApiProperties kimiApiProperties) {
        this.kimiApiProperties = kimiApiProperties;
    }

    @Override
    public String chat(String question) throws Exception {
        String requestBody = getRequestBody(question);

        Request request = new Request.Builder()
                .url(kimiApiProperties.getHostUrl())
                .post(RequestBody.create(requestBody, MediaType.parse("application/json; charset=utf-8")))
                .addHeader("Authorization", "Bearer " + kimiApiProperties.getApiKey())
                .build();
        Call call = new OkHttpClient().newCall(request);
        Response response = call.execute();

        String responseBody = response.body().string();
        if (!StringUtils.hasText(responseBody)) {
            return "无响应";
        }
        JSONObject responseObject = JSONObject.parseObject(responseBody);
        JSONArray choices = responseObject.getJSONArray("choices");
        StringBuilder result = new StringBuilder();
        for (Object choice : choices) {
            JSONObject oneChoice = (JSONObject) choice;
            JSONObject message = oneChoice.getJSONObject("message");
            String role = message.getString("role");
            String content = message.getString("content");
            if ("assistant".equalsIgnoreCase(role)) {
                result.append(content);
            }
        }
        return result.toString();
    }

    private String getRequestBody(String question) {
        JSONObject requestBody = new JSONObject();
        requestBody.put("model", kimiApiProperties.getModel());
        JSONArray messages = new JSONArray();
        JSONObject oneMessage = new JSONObject();
        oneMessage.put("role", "user");
        oneMessage.put("content", question);
        messages.add(oneMessage);
        requestBody.put("messages", messages);
        requestBody.put("temperature", 0.3);
        return requestBody.toJSONString();
    }
}
